package org.chnxi.datastructures.tree;

import java.io.*;
import java.util.*;

/**
 * 使用赫夫曼编码进行文件的压缩与解压
 */
public class HuffmanZipFile {

    /**
     * 解压文件
     * @param zipFile 待解压的文件
     * @param dstFile 解压后的文件
     */
    public static void unzipFile(String zipFile , String dstFile){
        //读取压缩文件内容
        byte[] zipFileData = null;
        Map<Byte , String> huffmanCodec = null;

        try(InputStream is = new FileInputStream(zipFile);
                ObjectInputStream ois =new ObjectInputStream(is)){
            zipFileData = (byte[]) ois.readObject();
            huffmanCodec = (Map<Byte,String>)ois.readObject();
        }catch (Exception e){
            throw new RuntimeException("读取压缩文件发生错误："+e.getMessage());
        }

        //使用赫夫曼编码进行解码
        byte[] unzipFileData = unzip(huffmanCodec , zipFileData);

        //将解码后内容写入文件
        try(OutputStream os = new FileOutputStream(dstFile);){
            os.write(unzipFileData);
            os.flush();
        }catch (Exception e){
            throw new RuntimeException("写入解压文件发生错误："+e.getMessage());
        }
    }

    /**
     * 压缩文件
     * @param srcFile 戴压缩的源文件
     * @param dstFile 压缩后的存储文件
     */
    public static void zipFile(String srcFile , String dstFile){
        //先获取文件内容
        byte[] srcFileData = null;

        try(InputStream is = new FileInputStream(srcFile)){
            srcFileData = new byte[is.available()];
            is.read(srcFileData);
        }catch (Exception e){
            throw new RuntimeException("读取源文件发生错误："+e.getMessage());
        }

        //统计出现次数
        List<HCNode> nodes = getHuffmanNodes(srcFileData);

        //创建赫夫曼树
        HCNode root = createHuffmanTree(nodes);

        //生成赫夫曼编码
        Map<Byte , String> huffmanCodes = getHuffmanCodes(root);

        //使用赫夫曼编码进行压缩
        byte[] huffmanCodeFileData = zip(srcFileData , huffmanCodes);

        //将压缩后的内容写入文件
        try(OutputStream os = new FileOutputStream(dstFile);
            ObjectOutputStream oos= new ObjectOutputStream(os)){
            oos.writeObject(huffmanCodeFileData);
            oos.writeObject(huffmanCodes);
            oos.flush();
        }catch (Exception e){
            throw new RuntimeException("关闭压缩文件写入流发生错误："+e.getMessage());
        }
    }

    /**
     * 将一个byte转换为一个二进制的字符串
     * @param flag ture：表示需要补高位，false：表示不需要补高位
     * @param data 传入的byte
     * @return 按照补码返回的二进制字符串
     */
    private static String byteToBitString(boolean flag, byte data){
        int tmp = data;
        if(flag){
            tmp |= 256;

        }
        String bitStr = Integer.toBinaryString(tmp);
        if(flag){
            return bitStr.substring(bitStr.length()-8);
        }else{
            return bitStr;
        }
    }

    /**
     * 完成解码
     * @param codes 赫夫曼编码表
     * @param datas 经过赫夫曼编码得到的字节数组
     * @return 原字符串对应的数组
     */
    private static byte[] unzip(Map<Byte , String> codes , byte[] datas){
        //获取二进制方式表示的数据字符串
        StringBuilder bitsStr = new StringBuilder();
        for (int i=0;i<datas.length;i++){
            boolean flag = i == (datas.length - 1);
            String byteStr = byteToBitString(!flag , datas[i]);
            bitsStr.append(byteStr);
        }

        //将赫夫曼编码表进行反转
        Map<String ,Byte> map = new HashMap<String, Byte>();
        for(Map.Entry<Byte , String> entry : codes.entrySet()){
            map.put(entry.getValue() , entry.getKey());
        }

        //对二进制字符串进行解码，将转换后的byte放入list中
        List<Byte> list = new ArrayList<Byte>();
        for (int i = 0; i < bitsStr.length(); ){ //在循环中自己控制i增长
            int count = 1; //控制i向后寻找长度
            boolean flag = true; //标记是否找到
            Byte b = null; //存储对应找到的byte

            while(flag){
                String key = bitsStr.substring(i, i+count); //i不动，逐渐递增向后寻找
                b = map.get(key);
                if(b == null){
                    count++;

                }else{
                    flag = false;
                }
            }

            list.add(b);
            i+=count;
        }

        //把list中的数据放入到byte[]并返回
        byte[] result = new byte[list.size()];
        for(int i=0;i<list.size();i++){
            result[i] = list.get(i);
        }

        return result;
    }

    /**
     * 将字符串对应的byte[]，通过生成的赫夫曼编码，返回一个赫夫曼编码压缩后的byte[]
     * @param bytes
     * @param huffmanCode
     * @return
     */
    private static byte[] zip(byte[] bytes , Map<Byte , String> huffmanCode){
        StringBuilder strBuilder = new StringBuilder();

        for(byte b : bytes){
            String value = huffmanCode.get(b);
            strBuilder.append(value);
        }

        int len = (strBuilder.length() + 7) /8;

        byte[] result = new byte[len];
        int resultIdx = 0;
        for (int i = 0; i < strBuilder.length(); i+=8){
            String strByte;

            if(i+8 > strBuilder.length()){
                strByte = strBuilder.substring(i);
            }else{
                strByte = strBuilder.substring(i , i+8);
            }

            result[resultIdx] = (byte) Integer.parseInt(strByte , 2);
            resultIdx++;
        }

        return result;
    }

    /**
     * 重载getCodes，方便调用
     * @param root
     * @return
     */
    private static Map<Byte , String> getHuffmanCodes(HCNode root){
        if(root == null){
            return null;
        }

        Map<Byte , String> huffmanCodes = new HashMap<Byte , String>();
        StringBuilder sb = new StringBuilder();

        getHuffmanCodes(root.left , "0" , sb , huffmanCodes);
        getHuffmanCodes(root.right , "1" , sb , huffmanCodes);
        return huffmanCodes;
    }

    /**
     * 生成赫夫曼编码
     * @param node 赫夫曼树节点
     * @param code 当前节点的编码取值
     * @param strBuilder 从root到当前节点的路径拼接
     * @param huffmanCodes 生成过的赫夫曼编码
     */
    private static void getHuffmanCodes(HCNode node, String code , StringBuilder strBuilder, Map<Byte , String> huffmanCodes){
        StringBuilder currStrBuilder = new StringBuilder(strBuilder);
        currStrBuilder.append(code);

        if(node != null){
            if(node.data == null){
                getHuffmanCodes(node.left , "0" , currStrBuilder , huffmanCodes);
                getHuffmanCodes(node.right , "1" , currStrBuilder , huffmanCodes);
            }else{
                huffmanCodes.put(node.data , currStrBuilder.toString());
            }
        }
    }

    /**
     * 创建赫夫曼树
     * @param nodes
     * @return
     */
    private static HCNode createHuffmanTree(List<HCNode> nodes){
        while (nodes.size() > 1){
            Collections.sort(nodes);

            HCNode left = nodes.get(0);
            HCNode right = nodes.get(1);
            HCNode parent = new HCNode(null , left.weiht+ right.weiht); //没有data内容
            parent.left = left;
            parent.right = right;

            nodes.remove(left);
            nodes.remove(right);
            nodes.add(parent);
        }

        return nodes.get(0);
    }

    /**
     * 接受字节数组，统计各个字符出现的对应次数
     * @param bytes
     * @return
     */
    private static List<HCNode> getHuffmanNodes(byte[] bytes){
        List<HCNode> nodes = new ArrayList<HCNode>();

        //遍历bytes，统计每一个byte出现的次数
        Map<Byte , Integer> counts = new HashMap<Byte , Integer>();
        for (byte b : bytes){
            Integer count = counts.get(b);
            if(count == null){
                counts.put(b , 1);
            }else {
                counts.put(b , count+1);
            }
        }

        // 把map中每个键值对创建为HCNode对象，放入nodes
        for(Map.Entry<Byte,Integer> entry : counts.entrySet()){
            nodes.add(new HCNode(entry.getKey() , entry.getValue()));
        }

        return nodes;
    }

    static class HCNode implements Comparable<HCNode> {
        Byte data; //存放数据本身
        int weiht; //权值，字符出现的次数
        HCNode left;
        HCNode right;

        protected HCNode(Byte data, int weight){
            this.data = data;
            this.weiht = weight;
        }

        @Override
        public int compareTo(HCNode o) {
            return this.weiht - o.weiht;
        }

        @Override
        public String toString() {
            return "HCNode{data=" + data + ", weiht=" + weiht + "}";
        }
    }

}
